[JAX] EP Dispatch with overflow detection option - #3277
Conversation
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( | ||
| cfg, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank | ||
| ) | ||
| out_leading = tuple(tokens.shape[:-1]) |
There was a problem hiding this comment.
Existing callers that unpack the documented two-value ep_prepare, four-value ep_dispatch, or two-value MoE results now receive an additional positional value unconditionally, causing ValueError: too many values to unpack even when overflow dropping is disabled.
Knowledge Base Used: JAX Fused Layers: Functions and Flax Modules
Greptile SummaryAdds optional expert-parallel overflow dropping and exposes each rank’s pre-drop receive demand.
Confidence Score: 4/5The PR is not yet safe to merge because its public EP and MoE functions still unconditionally change positional return arities and break existing unpacking callers. Existing callers written against the previous ep_prepare, ep_dispatch, moe, or Flax MoE tuple contracts receive additional values even when overflow dropping is disabled, producing runtime unpacking failures; the previously reported compatibility issue remains in the current code. Files Needing Attention: transformer_engine/jax/ep.py, transformer_engine/jax/moe.py, transformer_engine/jax/flax/moe.py Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant JAX as JAX EP/MoE API
participant FFI as JAX FFI
participant EP as Native EP communicator
User->>JAX: ep_bootstrap(drop_on_overflow)
JAX->>FFI: set bootstrap parameters
FFI->>EP: initialize communicator and buffers
User->>JAX: ep_prepare / ep_dispatch
JAX->>FFI: routing metadata and tensors
FFI->>EP: prepare and dispatch
EP-->>FFI: routed tensors and pre-drop demand
FFI-->>JAX: total_recv_tokens
JAX-->>User: outputs plus overflow signal
User->>JAX: ep_finalize()
JAX->>EP: release resources
Reviews (3): Last reviewed commit: "Merge branch 'main' into ep-jax-overflow" | Re-trigger Greptile |
phu0ngng
changed the title
| """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``.""" | ||
| """Exchange routing metadata for ``cfg``; return | ||
| ``(token_counts, total_recv_tokens, handle_mem)``. ``total_recv_tokens`` is | ||
| the per-rank pre-drop recv-slot total (includes tokens dropped on overflow).""" |
There was a problem hiding this comment.
When drop-on-overflow aka overflow-detection is disabled, this will now return a "trt" tensor that has an uninitialized value, right? Can we instead update this so ep_prepare will return (token_counts, None, handle_mem) instead when drop-on-overflow=False?
There was a problem hiding this comment.
No, it actually has the total_recv_tokens value i.e., the sum of token_counts.
The only difference here is that when drop-on-overflow=False (which is the default), as soon as there is an overflow, the kernel will trap and the program will crash.
There was a problem hiding this comment.
Oh I see, then the current implementation looks good. The PR LGTM once CI passes
|
/te-ci L1 JAX |
2 participants
Description
Adds a way to size EP recv buffers below the worst case and still stay correct.
ep_prepare/ep_dispatchnow exposetotal_recv_tokens, the per-rank pre-drop recv-slot demand, so callers can tell exactly when a step's routing exceedsrecv_capacity_per_rank; no need to conservatively provision every rank for the theoretical maximum.Bootstrap also gains an opt-in
drop_on_overflowpolicy so an occasional overflowing step is handled gracefully and keeps running (the default behavior is unchanged).Together these let users pick a tighter, cheaper recv capacity and monitor
total_recv_tokensto confirm it holds or to drive their own capacity policy.Type of change
Changes
ep_bootstrapgains adrop_on_overflowflag, plumbed throughset_ep_bootstrap_params(pybind) andNVTEEpGroupConfig.drop_on_overflow.ep_preparenow returns(token_counts, total_recv_tokens, handle_mem);ep_dispatch'scustom_vjpprimal correspondingly returnstotal_recv_tokensas a non-differentiable output.moe()and the Flax_MoEBlocknow return(output, aux_loss, total_recv_tokens); the extra output is non-differentiableep_finalize()(andreset_ep_config) to tear down the EPcommunicator so a process can re-
ep_bootstrapwith a different config.Checklist: